import os
import numpy as np
import json
import argparse
import open3d as o3d
from tqdm import tqdm
import sys

sys.path.append('../')  # 确保这个路径是正确的
from cyw_devkit.core.dataset import CywDataset
from my_io import copy_file
from tictoc import TicToc
import multiprocessing


def main(args):
    # 统计耗时
    cost = TicToc("标注数据生成")
    assert os.path.exists(args.data_path)
    frames = os.listdir(args.data_path)
    frames.sort(key=lambda x: x)
    files = []
    for dir in frames:
        if dir[:2] == '__':  # '__'
            dir = os.path.join(args.data_path, dir)
            if os.path.isdir(dir):
                files.append(dir)

    process_size = len(files)
    manager = multiprocessing.Manager()
    if process_size > 1:
        pool = multiprocessing.Pool(process_size)
        counter_list = manager.list()
        for idx in range(process_size):
            pool.apply_async(main_worker, args=(files[idx], args))
        pool.close()
        pool.join()
    else:
        main_worker(files[0], args)

    print("---------------------------------------------------------")
    print("处理完成: {}".format(files))
    cost.toc()
    print("---------------------------------------------------------")


def main_worker(dataset_path, args):
    if 'town' in dataset_path or 'carla' in dataset_path:
        radar_dim = 7
        normal_img=False
    else:
        radar_dim = 5
        normal_img=True

    cyw = CywDataset(dataset_path=dataset_path, radar_dim=radar_dim, normal_img= normal_img)
    tmp_dir = "anno_" + os.path.basename(dataset_path)
    tmp_anotate = os.path.join(args.anno_path, tmp_dir, 'lidar')
    label_dir = os.path.join(args.anno_path, tmp_dir, 'label')
    camera_id = 'camera75'
    tmp_anotate_img = os.path.join(args.anno_path, tmp_dir, 'camera', camera_id)
    tmp_anotate_calib = os.path.join(args.anno_path, tmp_dir, 'calib/camera')
    os.makedirs(tmp_anotate, exist_ok=True)
    os.makedirs(label_dir, exist_ok=True)
    os.makedirs(tmp_anotate_img, exist_ok=True)
    os.makedirs(tmp_anotate_calib, exist_ok=True)

    sensors_info = dict()
    sensors_info['extrinsic'] = cyw.data_dict['calib_info']['camera75']['transforms']['base_link']
    # sensors_info['extrinsic'] = np.array(sensors_info['extrinsic']).reshape(-1).tolist()
    sensors_info['extrinsic'] = np.linalg.inv(np.array(sensors_info['extrinsic'])).reshape(-1).tolist()
    # sensors_info['extrinsic'] = np.linalg.inv(np.array(sensors_info['extrinsic'])).tolist()
    sensors_info['intrinsic'] = cyw.data_dict['calib_info']['camera75']['normal']['K']
    sensors_info['intrinsic'] = np.array(sensors_info['intrinsic']).reshape(-1).tolist()
    json_path = os.path.join(tmp_anotate_calib, camera_id+'.json')
    with open(json_path, 'w') as f:
        json.dump(sensors_info, f, ensure_ascii=False, indent=2)
    for key_idx in tqdm(cyw.keyframe_idxs):
        key_frame = cyw.__getitem__(key_idx)
        key_stamp = key_frame['lidar'].header.stamp
        key_path = os.path.join(tmp_anotate, str(key_stamp) + '.pcd')
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(key_frame['lidar'].data[:, :3])
        o3d.io.write_point_cloud(key_path, pcd,
                                 write_ascii=True)  # save xyz.pcd
        # 拷贝图像和重命名，否则sus工具不能识别图像
        if 'town' in dataset_path or 'carla' in dataset_path:
            src_dir = os.path.join(cyw.dataset_path, 'camera', camera_id)
        else:
            src_dir = os.path.join(cyw.dataset_path, 'camera', camera_id+'_normal')
        file_name = str(key_frame['camera_dict'][camera_id].header.stamp) + '.jpg'
        new_name = str(key_stamp) + '.jpg'
        dst_dir = tmp_anotate_img
        copy_file(os.path.join(src_dir, file_name), dst_dir)

        os.rename(os.path.join(dst_dir, file_name), os.path.join(dst_dir, new_name))

        if args.motion:
            label_path = os.path.join(cyw.dataset_path,'samples/label_e2e', str(key_stamp) + '.json')
            copy_file(label_path, label_dir)

        pass


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Configuration Parameters')
    parser.add_argument('--data-path', default="/media/adt/T7/ZWH/docker/files/data/motion/datasets",
                        help='your data root')
    parser.add_argument('--anno-path', default="/media/adt/T7/ZWH/docker/files/data/motion/annotate",
                        help='')
    parser.add_argument('--radar-dim', type=int, default=5,
                        help='')
    parser.add_argument('--motion', type=bool, default=True,
                        help='')
    args = parser.parse_args()

    main(args)
